import asyncio
import websockets
import json

# 用于存储客户端WebSocket连接的字典
clients = {}

async def signaling_server(websocket, path):
    # 当新的WebSocket连接被打开时，将其添加到客户端字典中
    clients[websocket] = None
    try:
        # 接收消息直到建立WebRTC连接
        async for message in websocket:
            # 尝试解析JSON格式的消息
            try:
                message_dict = json.loads(message)
            except json.JSONDecodeError:
                print(f"Received invalid JSON from {websocket.remote_address}: {message}")
                continue

            # 打印接收到的消息
            print(f"Received message from {websocket.remote_address}: {message_dict}")
            if message_dict.get('type', None) == 'p2p':
                for client in clients:
                    if client != websocket:
                        await client.send(json.dumps(message_dict))

            # 如果是offer或answer，广播给所有其他客户端
            if 'type' in message_dict and message_dict['type'] in ['offer', 'answer']:
                for client in clients:
                    if client != websocket:
                        await client.send(json.dumps(message_dict))
            # 如果是ICE候选，只发送给对应的客户端
            elif 'candidate' in message_dict and 'peer_id' in message_dict:
                peer_id = message_dict['peer_id']
                for client in clients:
                    if clients[client] == peer_id:
                        await client.send(json.dumps(message_dict))
                        break
    except websockets.exceptions.ConnectionClosed:
        # 当连接关闭时，从字典中移除客户端
        del clients[websocket]
    finally:
        print(f"WebSocket connection closed: {websocket.remote_address}")

# 启动WebSocket服务端
async def main():
    async with websockets.serve(signaling_server, "localhost", 8765):
        await asyncio.Future()  # 运行直到被关闭

if __name__ == "__main__":
    asyncio.run(main())
